#!/usr/bin/python
from rap_core import *


zeff=10
core=-10.
ip=0.5


clevs=rap_states_ca()
clevs.set_nuclear_charge(zeff)
clevs.set_core_energy(core)
clevs.set_ionization_threshold(ip)

clevs.add_level('1s2',0.0)
clevs.add_level('1s 2s',0.1)
clevs.add_level('1s 2p',0.11)
clevs.add_level('1s 3s',0.2)
clevs.add_level('1s 3p',0.21)
clevs.add_level('1s 3d',0.22)

if clevs.nuclear_charge() != zeff:
  print 'problem with nuclear_charge()'
if clevs.core_energy() != core:
  print 'problem with core_energy()'
if clevs.ionization_threshold() != ip:
  print 'problem with ionization_threshold()'
if len(clevs) != 6:
  print 'wrong number of levels'

if str(clevs[1]) != '1s2':
  print 'A - incorrect level: 1s2'
if str(clevs[6]) != '1s1 3d1':
  print 'A - incorrect level: 1s 3d'


toadd=[]
toadd.append((0.00,2,'1s2'))
toadd.append((0.20,3,'1s 2s'))
toadd.append((0.21,4,'1s 2p'))
toadd.append((0.30,5,'1s 3s'))
toadd.append((0.31,6,'1s 3p'))
toadd.append((0.32,7,'1s 3d'))

clevs2=clevs.copy()
clevs2.clear_levels()
clevs2.add_levels(toadd)

if clevs2.nuclear_charge() != zeff:
  print 'problem with nuclear_charge()'
if clevs2.core_energy() != core:
  print 'problem with core_energy()'
if clevs2.ionization_threshold() != ip:
  print 'problem with ionization_threshold()'
if len(clevs2) != 6:
  print 'wrong number of levels'

if str(clevs2[2]) != '1s2':
  print 'B - incorrect level: 1s2'
if str(clevs2[7]) != '1s1 3d1':
  print 'B - incorrect level: 1s 3d'

clevs.set_lmap(clevs2)
if clevs.label(3,imap=True,spacer='.') != '1s1.2s1':
  print 'problem with mapping or label'

clevs2.reindex()
if str(clevs2[1]) != '1s2':
  print 'C - incorrect level: 1s2'
if str(clevs2[6]) != '1s1 3d1':
  print 'C - incorrect level: 1s 3d'




